import xarray as xr
import numpy as np
from scipy.interpolate import griddata
import netCDF4 as nc
import pandas as pd
import multiprocessing

# 读取FY4A坐标（只需读取一次）
coord_file = nc.Dataset('/home/liudd/data_preprocessing/FY4A_coordinates.nc', 'r')
lon_fy4a = coord_file.variables['lon'][:, :].T
lat_fy4a = coord_file.variables['lat'][:, :].T
lon_fy4a_360 = lon_fy4a % 360

# 定义通用插值函数
def interpolate_to_fy4a(data_var, lon_era5, lat_era5):
    # 生成ERA5网格点
    lon_grid, lat_grid = np.meshgrid(lon_era5, lat_era5)
    points = np.column_stack((lon_grid.ravel(), lat_grid.ravel()))

    # 准备目标点
    target = np.column_stack((lon_fy4a_360.ravel(), lat_fy4a.ravel()))
    valid = ~np.isnan(target).any(axis=1)

    # 执行插值
    result = np.full(lon_fy4a.shape, np.nan)
    interpolated = griddata(
        points,
        data_var.ravel(),
        target[valid],
        method='linear',
        fill_value=np.nan
    )
    result.ravel()[valid] = interpolated
    return result

# 处理单个时间点数据的函数
def process_single_timepoint(i, sp_data, t2m_data, q_data, lon_era5, lat_era5, times):
    print(f"Processing time index {i}...")

    # 提取当前时间点的数据
    current_sp = sp_data[i]
    current_t2m = t2m_data[i]
    current_q = q_data[i]

    # 处理比湿数据
    surface_q = q_data[:, 0, :, :][i]  # 假设 pressure_level=1000 是第一个维度

    # 执行插值
    t2m_interp = interpolate_to_fy4a(current_t2m, lon_era5, lat_era5)
    sp_interp = interpolate_to_fy4a(current_sp, lon_era5, lat_era5)
    q_interp = interpolate_to_fy4a(surface_q, lon_era5, lat_era5)

    # 创建当前时间点的数据集
    ds = xr.Dataset(
        {
            "temp_2m": (("y", "x"), t2m_interp),
            "surface_pressure": (("y", "x"), sp_interp),
            "surface_specific_humidity": (("y", "x"), q_interp)
        },
        coords={
            "time": times[i],
            "lon": (("y", "x"), lon_fy4a),
            "lat": (("y", "x"), lat_fy4a)
        }
    )

    # 添加属性
    ds['time'].attrs['long_name'] = 'Time'
    ds['lon'].attrs['units'] = 'degrees_east'
    ds['lat'].attrs['units'] = 'degrees_north'

    # 生成文件名
    time_str = pd.to_datetime(times[i]).strftime('%Y%m%d%H')
    filename = f'/mnt/datastore/liudddata/ERA5_output/202001_04/ERA5_{time_str}_surface_vars_FY4A_grid.nc'

    # 保存结果
    ds.to_netcdf(filename)
    print(f"{time_str}数据处理完成并已保存到 {filename}。")
    return filename

# 主函数
def main():
    # 读取包含多个时间点数据的地表数据文件
    era5_surface = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200104/temp20200201-29.nc')
    sp = era5_surface['sp'].values
    t2m = era5_surface['t2m'].values
    lon_era5 = era5_surface.longitude.values
    lat_era5 = era5_surface.latitude.values
    times = era5_surface.valid_time.values

    # 读取包含多个时间点数据的比湿数据文件
    era5_q = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200104/specifichumidity_1000hPa20200102_29.nc')
    q = era5_q['q'].values
    levels = era5_q.pressure_level.values.astype(float)

    # 根据数据量动态调整进程数
    cpu_count = multiprocessing.cpu_count()
    safe_process_num = max(1, int(cpu_count * 0.8))  # 使用80%的CPU资源

    # 创建进程池
    with multiprocessing.Pool(processes=safe_process_num) as pool:
        # 准备参数
        args = [(i, sp, t2m, q, lon_era5, lat_era5, times) for i in range(len(times))]

        # 使用imap_unordered提高效率
        for result in pool.starmap(process_single_timepoint, args):
            pass

    print("所有逐小时数据处理完成。")

if __name__ == "__main__":
    main()